from django.db import migrations

from vulnerabilities.severity_systems import SCORING_SYSTEMS

"""
This migration fixes the CVSS severity scores.

The starting point is that we have two VulnerabilitySeverity records for
each CVSS scores:

- one record with numeric score value and a "cvssv2/v3" system
- a second record with a CVSS vector and a "cvssv2/v3_vector" system;

The migration merges the two records where the CVSS vector becomes the scoring
element of the score value, and removes the extra record.

In the odd cases where we have no a pair of records we use a best effort to fix
the data.
"""


def merge_cvss_score_and_vector_severities(apps, _):
    """
    Merge CVSS score and vector VulnerabilitySeverity and remove the
    cvssv*_vector severities after merging.
    """
    cvss_systems = [
        "cvssv2",
        "cvssv3",
        "cvssv3.1",
    ]

    cvss_vector_systems = [
        "cvssv2_vector",
        "cvssv3_vector",
        "cvssv3.1_vector",
    ]

    all_cvss_systems = cvss_systems + cvss_vector_systems

    VulnerabilitySeverity = apps.get_model("vulnerabilities", "VulnerabilitySeverity")

    updated_severities_to_save = {}
    redundant_severity_ids_to_delete = set()

    for severity in VulnerabilitySeverity.objects.filter(scoring_system__in=all_cvss_systems):
        process_severity(severity, updated_severities_to_save, redundant_severity_ids_to_delete)

    # finally batch update and delete

    # batch delete the collected vecto severity ids
    deleted, _ = VulnerabilitySeverity.objects.filter(
        id__in=redundant_severity_ids_to_delete
    ).delete()
    print(f"Deleted {deleted} CVSS VulnerabilitySeverity")

    # update in bulk, 1000 at a time
    updated = VulnerabilitySeverity.objects.bulk_update(
        objs=updated_severities_to_save.values(),
        fields=[
            "scoring_system",
            "scoring_elements",
            "value",
        ],
        batch_size=1000,
    )
    print(f"Updated {updated} CVSS VulnerabilitySeverity")

    leftover_vectors = VulnerabilitySeverity.objects.filter(
        scoring_system__in=cvss_vector_systems
    ).count()
    if leftover_vectors:
        print(f"ERRROR!!!! {leftover_vectors} CVSS vector VulnerabilitySeverity left")
        for leftover in VulnerabilitySeverity.objects.filter(
            scoring_system__in=cvss_vector_systems
        ):
            print(leftover)


def process_severity(severity, updated_severities_to_save, redundant_severity_ids_to_delete):
    """
    Add ``severity`` to the ``updated_severities_to_save`` mapping and update the
    ``redundant_severity_ids_to_delete`` set of redundant severities as needed.
    If ``severity`` already exists keep one record merging the data
    and treat as deletable the other record.
    """
    if not severity.value or not severity.value.strip():
        redundant_severity_ids_to_delete.add(severity.pk)
        return

    # convert scoring system to non vector
    # >>> "cvssv2_vector".partition("_vector")
    # ('cvssv2', '_vector', '')
    # >>> "cvssv2".partition("_vector")
    # ('cvssv2', '', '')
    scoring_system, is_vector, _ = severity.scoring_system.partition("_vector")
    if is_vector:
        # compute score and move vector to scoring_elements
        severity.scoring_system = scoring_system
        vector = severity.value
        severity.scoring_elements = vector
        severity.value = SCORING_SYSTEMS[scoring_system].compute(vector)

    # Build a unique hashable key for a severity to keep track of duplicates
    # The VulnerabilitySeverity model unique_together is:
    # ["reference", "scoring_system", "value"]
    key = (
        severity.reference.pk,
        severity.scoring_system,
        severity.value,
    )

    existing = updated_severities_to_save.get(key)

    if not existing:
        # keep instance for bulk update
        updated_severities_to_save[key] = severity

    else:
        # merge/update existing with vector if
        if not existing.scoring_elements and severity.scoring_elements:
            existing.scoring_elements = severity.scoring_elements
        # keep id for bulk deletion
        redundant_severity_ids_to_delete.add(severity.pk)


class Migration(migrations.Migration):
    dependencies = [
        ("vulnerabilities", "0031_vulnerabilityseverity_scoring_elements"),
    ]

    operations = [
        migrations.RunPython(merge_cvss_score_and_vector_severities, migrations.RunPython.noop),
    ]
